#include "AIChatbot.h"


AIChatbot::AIChatbot()
    : _apiKey(""), _selectedAI(""), _selectedAIVersion("") {
    addHistoryMessage("system", "");
}

AIChatbot::AIChatbot(const String& key, const String& ai) {
    setKey(key);
    selectAI(ai);
    addHistoryMessage("system", "");
}

AIChatbot::AIChatbot(const String& key, const String& ai, const String& message) {
    setKey(key);
    selectAI(ai);
    addHistoryMessage("system", message);
}


bool AIChatbot::validateKeys() {
    if (!_url.length()) {
        _info = F("[Error]: LLM API Provider is not set.");
        return false;
    }

    if (!_apiKey.length()) {
        _info = F("[Error]: API key is not set.");
        return false;
    }

    if (!_selectedAIVersion.length()) {
        _info = F("[Error]: Model is not set.");
        return false;
    }

    return true;
}

void AIChatbot::setKey(const String& key) {
    _apiKey = key;
}

void AIChatbot::setTimeout(uint16_t timeout) {
    _timeout = timeout;
}

void AIChatbot::setMaxRetries(uint8_t maxRetries) {
    _maxRetries = maxRetries;
}

void AIChatbot::setSystemPrompt(const String& message) {
    _messages[0]["content"] = message;
}

void AIChatbot::setCustomUrl(const String& url) {
    _url = url;
}

void AIChatbot::selectAI(const String& ai) {
    uint8_t i = ai.indexOf('@');
    if (i == -1) {
        _selectedAI = "";
        _selectedAIVersion = "";
        return;
    }
    String aiName = ai.substring(0, i);
    String aiVersion = ai.substring(i + 1, ai.length());
    selectAI(aiName, aiVersion);
}

void AIChatbot::selectAI(const String& aiName, const String& aiVersion) {
    _selectedAI = aiName;
    if (_selectedAI == "ChatGPT") {
        _url = AI_CHATBOT_CHAT_GPT_URL;
    } else if (_selectedAI == "HuggingFace") {
        _url = AI_CHATBOT_HUGGING_FACE_URL;
    } else if (_selectedAI == "DeepSeek") {
        _url = AI_CHATBOT_DEEP_SEEK_URL;
    } else if (_selectedAI == "SiliconFlow") {
        _url = AI_CHATBOT_SILICON_FLOW_URL;
    } else if (_selectedAI == "VolcEngine") {
        _url = AI_CHATBOT_VOLC_ENGINE_URL;
    } else if (_selectedAI == "LingYiWanWu") {
        _url = AI_CHATBOT_LING_YI_WAN_WU_URL;
    } else if (_selectedAI == "MoonShot") {
        _url = AI_CHATBOT_MOON_SHOT_URL;
    } else if (_selectedAI == "BaiChuanAI") {
        _url = AI_CHATBOT_BAI_CHUAN_AI_URL;
    } else if (_selectedAI == "TencentHunYuan") {
        _url = AI_CHATBOT_TENCENT_HUN_YUAN_URL;
    } else if (_selectedAI == "AliyunBaiLian") {
        _url = AI_CHATBOT_ALIYUN_BAI_LIAN_URL;
    }
    _selectedAIVersion = aiVersion;
}

void AIChatbot::enableHistoryMessages(bool state) {
    _history = state;
}

void AIChatbot::setHistoryMessagesNum(uint8_t num) {
    _maxHistoryNum = num;
}

void AIChatbot::emptyHistoryMessages() {
    _messages.clear();
}

void AIChatbot::addHistoryMessage(const String& role, const String& content) {
    JsonDocument doc;
    JsonObject object = doc.to<JsonObject>();
    object["role"] = role;
    object["content"] = content;
    _messages.add(object);
}

void AIChatbot::clearUserHistoryMessages() {
    _messages.clear();
    uint8_t len = _messages.size();
    for (uint8_t i = 1; i < len; i++) {
        _messages.remove(i);
    }
}

String AIChatbot::chat(const String& message, uint16_t maxTokens, void (*contentCallback)(String), void (*reasoningContentCallback)(String)) {
    if (!validateKeys()) {
        if (contentCallback != NULL) {
            contentCallback(_info);
        }
        return _info;
    }
    String response;
    if (_history) {
        addHistoryMessage("user", message);
    } else {
        if (_messages.size() != 2) {
            clearUserHistoryMessages();
            addHistoryMessage("user", message);
        } else {
            if (_messages[1]["role"] != "user") {
                _messages[1]["role"] = "user";
            }
            _messages[1]["content"] = message;
        }
    }
    String serializedMessages;
    serializeJson(_doc, serializedMessages);
    if (_selectedAI == "HuggingFace") {
        response = sendToHuggingFace(serializedMessages, maxTokens, contentCallback, reasoningContentCallback);
    } else {
        response = sendToCommonGPT(_url, serializedMessages, maxTokens, contentCallback, reasoningContentCallback);
    }
    if (_history) {
        addHistoryMessage("assistant", response);
        uint8_t messages_len = _messages.size();
        if (_maxHistoryNum > 0) {
            int len = (messages_len - 1) - _maxHistoryNum * 2;
            for (uint8_t i = 0; i < len; i++) {
                _messages.remove(1);
            }
        }
    }
    return response;
}

String AIChatbot::sendToCommonGPT(const String& url, const String& message, uint16_t maxTokens, void (*contentCallback)(String), void (*reasoningContentCallback)(String)) {
    String payload = "{\"model\": \"" + _selectedAIVersion + 
        "\", \"messages\": " + message + ", \"stream\": " + (contentCallback ? "true" : "false") + ", \"max_tokens\": " + maxTokens + "}";
    return makeHttpRequest(url, payload, _apiKey, contentCallback, reasoningContentCallback);
}

String AIChatbot::sendToHuggingFace(const String& message, uint16_t maxTokens, void (*contentCallback)(String), void (*reasoningContentCallback)(String)) {
    String payload = "{\"inputs\": \"" + message + "\"}";
    return makeHttpRequest(AI_CHATBOT_HUGGING_FACE_URL + _selectedAIVersion, payload, _apiKey, contentCallback, reasoningContentCallback);
}

String AIChatbot::makeHttpRequest(const String& url, const String& payload, const String& apiKey, void (*contentCallback)(String), void (*reasoningContentCallback)(String)) {
    HTTPClient httpClient;
    WiFiClientSecure wifiClient;
    String response;
    wifiClient.setInsecure();

    if (!httpClient.begin(wifiClient, url)) {
        response = F("[Error]: HTTP begin failed.");
        if (contentCallback != NULL) {
            contentCallback(response);
        }
        return response;
    }

    httpClient.setTimeout(_timeout);
    httpClient.addHeader("Content-Type", "application/json");
    if (contentCallback != NULL) {
        httpClient.addHeader("Accept", "text/event-stream");
    } else {
        httpClient.addHeader("Accept", "application/json");
    }
    httpClient.setReuse(true);
    httpClient.addHeader("Authorization", "Bearer " + apiKey);
    uint8_t retryCount = 0;
    int httpResponseCode = -1;

    while (retryCount < _maxRetries) {
        httpResponseCode = httpClient.POST(payload);
        if (httpResponseCode > 0) {
            break;
        } else {
            retryCount++;
            delay(1000);
        }
    }

    if (httpResponseCode != HTTP_CODE_OK) {
        if (httpResponseCode < 0) {
            response = F("[Error]: Network disconnected.");
        } else {
            response = httpClient.getString();
            httpClient.end();
        }
        if (contentCallback != NULL) {
            contentCallback(response);
        }
        return response;
    }

    if (contentCallback == NULL) {
        String line = httpClient.getString();
        JsonDocument doc;
        DeserializationError error = deserializeJson(doc, line);
        if (!error && doc["choices"][0]["message"]["content"].is<String>()) {
            response = doc["choices"][0]["message"]["content"].as<String>();
        }
        httpClient.end();
        return response;
    }

    while (wifiClient.connected()) {
        if (wifiClient.available()) {
            String line = wifiClient.readStringUntil('\n');
            if (!line.startsWith("data:")) {
                continue;
            }
            if (line.endsWith("[DONE]")) {
                break;
            }
            JsonDocument doc;
            line = line.substring(6, line.length());
            DeserializationError error = deserializeJson(doc, line);
            if (error) {
                continue;
            }
            JsonVariant variant = doc["choices"][0]["delta"].as<JsonVariant>();
            if (variant["reasoning_content"].is<String>() && reasoningContentCallback != NULL) {
                String tokens = variant["reasoning_content"].as<String>();
                reasoningContentCallback(tokens);
            }
            if (variant["content"].is<String>()) {
                String tokens = variant["content"].as<String>();;
                response += tokens;
                contentCallback(tokens);
            }
        }
    }
    httpClient.end();
    return response;
}